import numpy as np
import networkx as nx
import matplotlib
import matplotlib.pyplot as plt
# 设置字体为 SimHei
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False


def LMSTCLU_OPT(A, dist, nc, minsize, clu_num):
    """
    在局部密度峰值上构建 MST 并做子簇分裂
    输入:
        A: numpy array, shape (M, dim) — 核点位置
        dist: ndarray, shape (M, M) — 距离矩阵
        nc: list or array of length M — 每个核点对应的簇大小
        minsize: int — 子簇最小尺寸阈值
        clu_num: int — 期望总簇数
    输出:
        cl: ndarray, shape (M,) — 每个核点的簇标签
    """


    # 1. 构造最小生成树
    M = A.shape[0]
    ST = construct_mst(dist, method='prim')  # 使用 Prim 算法
    ST = ST + ST.T  # 对称化

    # 绘制最小生成树
#    plot_mst(A, ST)

    # 2. 提取所有边并按权重降序
    idxs = np.where(np.triu(ST) > 0)    # 找到最小生成树邻接矩阵 ST 中所有非零元素的位置
    edges = ST[idxs]    # 提取边的权重
    sorted_inds = np.argsort(edges)[::-1]   # 按权重降序排列
    edge_pairs = [(idxs[0][i], idxs[1][i]) for i in sorted_inds] # 根据排序后的边权重索引，生成最小生成树中所有边的节点对列表

    # 辅助：BFS 计算带权大小
    def bfs_weight(start, matrix):
        visited = np.zeros(M, dtype=bool)
        queue = [start]
        visited[start] = True
        total = 0
        while queue:
            u = queue.pop(0)
            total += nc[u]
            for v in range(M):
                if matrix[u, v] > 0 and not visited[v]:
                    visited[v] = True
                    queue.append(v)
        return total

    # 3. 逐条“剪”边，直到簇数达到 clu_num
    ST2 = ST.copy()
    k = 1
    for p, q in edge_pairs:
        if k >= clu_num:
            break
        if ST2[p, q] == 0:
            continue
        # 模拟移除 p-q 边
        tmp = ST2.copy()
        tmp[p, q] = 0
        tmp[q, p] = 0
        s1 = bfs_weight(p, tmp)
        s2 = bfs_weight(q, tmp)
        if s1 >= minsize and s2 >= minsize:
            ST2 = tmp
            k += 1

    # 4. 最终连通分量标记
    cl = np.zeros(M, dtype=int)
    visited = np.zeros(M, dtype=bool)
    label = 0
    for i in range(M):
        if not visited[i]:
            label += 1
            queue = [i]
            visited[i] = True
            cl[i] = label
            while queue:
                u = queue.pop(0)
                for v in range(M):
                    if ST2[u, v] > 0 and not visited[v]:
                        visited[v] = True
                        cl[v] = label
                        queue.append(v)
    return cl



def construct_mst(dist, method='kruskal'):
    """
    构造最小生成树 (MST)
    输入:
        dist: ndarray, shape (M, M) — 距离矩阵
        method: str — MST 构造方法 ('kruskal' 或 'prim')
    输出:
        ST: ndarray, shape (M, M) — 最小生成树的邻接矩阵
    """
    # 将距离矩阵转换为图
    G = nx.Graph()
    M = dist.shape[0]
    for i in range(M):
        for j in range(i + 1, M):  # 避免重复边
            if dist[i, j] > 0:
                G.add_edge(i, j, weight=dist[i, j])

    # 构造 MST
    if method == 'kruskal':
        mst = nx.minimum_spanning_tree(G, algorithm='kruskal')
    elif method == 'prim':
        mst = nx.minimum_spanning_tree(G, algorithm='prim')
    else:
        raise ValueError("Unsupported method. Use 'kruskal' or 'prim'.")

    # 转换为邻接矩阵
    ST = nx.to_numpy_array(mst)
    return ST

def plot_mst(A, pred):
    """
    Draws the minimum spanning tree (MST).
    A: ndarray, shape (N, 2) — Data point coordinates
    pred: ndarray, shape (N, N) — Adjacency matrix of the MST
    """
    print("绘制最小生成树...")
    N = A.shape[0]
    plt.figure()
    plt.scatter(A[:, 0], A[:, 1], c='k', marker='.', label='Data Points')  # Plot data points
    for i in range(N):
        for j in range(N):
            if pred[i, j] > 0:  # Check if there is an edge between nodes i and j
                x = [A[i, 0], A[j, 0]]
                y = [A[i, 1], A[j, 1]]
                plt.plot(x, y, linewidth=3, color='g')  # Plot the edge
    plt.title("Minimum Spanning Tree")
    plt.legend()
    plt.show()